# pyoffwiki - offline wikipedia browser for the Irex Iliad
# written by Amir Shimoni
# put into the public domain for any use
#   but please give credits to those who deserve them...

import sys
import os
import zlib
import re
import pycdb
import seekbzip2
import struct
from xml.sax.saxutils import unescape
from mwlib import cdb

try:
    from xml.etree import cElementTree
except ImportError:
    import cElementTree

ns = '{http://www.mediawiki.org/xml/export-0.3/}'

wikiindex = "wikiidx"
wikidata = "wikidata.bin"



def normname(name):
    name = name.strip().replace("_", " ")
    name = name.lower()
    #name = name[:1].upper()+name[1:].lower()
    return name

class Bz2WikiDB(object):
    redirect_rex = re.compile(r'^#Redirect:?\s*?\[\[(?P<redirect>.*?)\]\]', re.IGNORECASE)
    templates_enabled=False

    def __init__(self, mainbz2file):
        basefn = mainbz2file[:len(mainbz2file)-4]
        #self.cdb_file = open(basefn+'.cdb', 'rb')
        self.cdb_file = seekbzip2.SeekableBzip2File(basefn+'.cdb.bz2',basefn+'.cdb.bzt')
        self._thewiki = seekbzip2.SeekableBzip2File(basefn+".bz2", basefn+".bzt")

        #self.dir = dir
        #self.obj2pos_path = os.path.join(self.dir, wikidata)
        #self.cdb = cdb.Cdb(open(os.path.join(self.dir, wikiindex+'.cdb'), 'rb'))
        self.isThereMore = None

    def getRandomTitle(self):
        return pycdb.cdbrandom(self.cdb_file)

    def getCdbListAt(self, startpos, size):
        return pycdb.cdbGetKeyList(self.cdb_file, startpos, size)

    def getCdbSearch(self, key):
        return pycdb.cdbSearch(self.cdb_file, key)

    def setTemplates(self, isEnabled):
        self.templates_enabled = isEnabled

    def isMoreArticles(self, title_):
        return self.isThereMore == title_

    def getParsedArticle(self, title, uparser):
        rd = self.getRawArticle(title)
        #rd = unescape(rd.decode('utf-8'))
        #rd = rd.decode('utf-8')
        #rd = unicode(rd, 'utf-8')
        #rd = rd.encode('utf-8', '?')
        #print rd
        return uparser.parseString(title, raw=rd, wikidb=self)

    def getRawArticle(self, title, raw=None, revision=None):
        origtitle = title
        print "Getting article:", title

        title = normname(title)
        try:
            tt = pycdb.cdbget(self.cdb_file, 0, title)
        except KeyError:
            print "Not found:", title
            return None

        try:
            if pycdb.cdbget(self.cdb_file, 0, title+"."):
                print "Another match:",title+"."
                self.isThereMore = origtitle
        except KeyError:
            None
            
        tp = struct.unpack('<QL', tt[0])
        self._thewiki.seek(tp[0])
        thetext = unicode(unescape(self._thewiki.read(tp[1])), 'utf-8')
        #thetext = unescape(self._thewiki.read(tp[1]))
        #thetext = self._thewiki.read(tp[1])

        if False:
            mo = self.redirect_rex.search(thetext)
            if mo:
                redirect = mo.group('redirect')
                redirect = normname(redirect.split("|", 1)[0].split("#", 1)[0])
                
                return self.getRawArticle(redirect)

        return thetext

    def getTemplate(self, title, followRedirects=False):
        if not self.templates_enabled:
            return None
        if not title.startswith("Template:"):
            title = "Template:" + title

        title = normname(title)
        try:
            tt = pycdb.cdbget(self.cdb_file, 0, title)
        except KeyError:
            print "Not found:", title
            return None
        tp = struct.unpack('<QL', tt[0])
        self._thewiki.seek(tp[0])
        thetext = unicode(unescape(self._thewiki.read(tp[1])), 'utf-8')
        return thetext

        if ":" in title:
            title = title.split(':', 1)[1]

        title = normname(title)
        res = unicode(self._readobj(u"T:"+title) or "", 'utf-8')
        if not res:
            return res

        mo = self.redirect_rex.search(res)
        if mo:
            redirect = mo.group('redirect')
            redirect = normname(redirect.split("|", 1)[0].split("#", 1)[0])
            return self.getTemplate(redirect)
        return res


#    def articles(self):
#        for k, v in self.cdb:
#            if k[0]==':':
#                k = unicode(k[1:], "utf-8")
#                yield k
